package com.sz.szaiagent.rag;

import org.springframework.ai.document.Document;
import org.springframework.ai.transformer.splitter.TokenTextSplitter;
import org.springframework.stereotype.Component;

import java.util.List;

/**
 * 自定义基于 Token 的切词器
 * 该类用于将文档按Token进行分割，便于后续处理和分析。
 */
@Component
class MyTokenTextSplitter {
    /**
     * 使用默认配置的TokenTextSplitter对文档进行分割
     * @param documents 待分割的文档列表
     * @return 分割后的文档列表
     */
    public List<Document> splitDocuments(List<Document> documents) {
        TokenTextSplitter splitter = new TokenTextSplitter();
        return splitter.apply(documents); // 应用默认配置的TokenTextSplitter进行文档分割
    }

    /**
     * 使用自定义配置的TokenTextSplitter对文档进行分割
     * @param documents 待分割的文档列表
     * @return 分割后的文档列表
     */
    public List<Document> splitCustomized(List<Document> documents) {
        TokenTextSplitter splitter =
                new TokenTextSplitter(200, 100, 10, 5000, true); // 自定义TokenTextSplitter配置
        return splitter.apply(documents); // 应用自定义配置的TokenTextSplitter进行文档分割
    }
}
